use rand::prelude::*;
use rand_distr::{Distribution, Normal};
use rayon::prelude::*;
use std::f64::consts::PI;
use std::fs::File;
use std::io::prelude::*;
use std::io::Write;
use std::time::SystemTime;

#[derive(Clone, Debug)]
pub struct Quantizer {
    bit_count: u32,
    full_scale: f64,
    one_bit: f64,
}

impl Quantizer {
    pub fn new(bit_count: u32, full_scale: f64) -> Quantizer {
        Quantizer {
            bit_count,
            full_scale,
            one_bit: 2.0 / ((1 << (bit_count - 1)) as f64),
        }
    }
    fn quantify(&self, d: f64) -> f64 {
        if d > self.full_scale {
            self.full_scale
        } else if d < -self.full_scale {
            -self.full_scale
        } else {
            (d / self.one_bit).round() * self.one_bit
        }
    }
}

#[derive(Clone, Debug)]
struct BBPLLState {
    rng_g: Normal<f64>,
    rng_ref: Normal<f64>,
    jitter: f64,
    jitter_ref: f64,
    n: u32,
    g: f64,
    g_charge_accelerating: f64,
    lms_gain: f64,
    lms_gain_ca: f64,
    alpha_counter: f64,
    delta_sigma_counter: f64,
    delta_sigma_counter_ca: f64,
    g_dac_vco: f64,
    dac_q: Quantizer,
    ca_q: Quantizer,
    lms_counter: f64,
    lms_counter_ca: f64,

    std_period_counter: u64,
    std_phase: f64,
    ref_period_counter: u64,
    ref_phase: f64,
    ref_jitter_phase: Vec<f64>,
    dco_period_counter: u64,
    dco_phase: f64,
    dco_jitter_phase: Vec<f64>,
    dco_jitter_record: Vec<f64>,
    dco_time_record: Vec<f64>,
    ref_jitter_record: Vec<f64>,
    ref_time_record: Vec<f64>,
    debug_record: Vec<f64>,
    gain_record: Vec<f64>,
    ref_jitter_acc: f64,
    dco_control_word: f64,
    dco_cw_raw: f64,
    ca_control_word: f64,
    ca_cw_raw: f64,

    sampler: f64,
    rawsampler: f64,
    last_sample: Vec<f64>,
    delay_n: usize,
}

impl BBPLLState {
    fn new(
        jitter: f64,
        jitter_ref: f64,
        N: u32,
        g: f64,
        g_charge_accelerating: f64,
        lms_gain: f64,
        lms_gain_ca: f64,
        g_dac_vco: f64,
        dac_q_bit_count: u32,
        dac_q_full_scale: f64,
        ca_q_bit_count: u32,
        ca_q_full_scale: f64,
        delay_n: usize,
    ) -> BBPLLState {
        BBPLLState {
            rng_g: Normal::new(0.0, std::f64::consts::SQRT_2 * jitter).unwrap(),
            rng_ref: Normal::new(0.0, std::f64::consts::SQRT_2 * jitter_ref).unwrap(),
            jitter,
            jitter_ref,
            n: N,
            g,
            g_charge_accelerating,
            lms_gain,
            lms_gain_ca,
            alpha_counter: 0.0,
            delta_sigma_counter: 0.0,
            delta_sigma_counter_ca: 0.0,
            lms_counter: 0.0,
            lms_counter_ca: 0.0,
            g_dac_vco,
            dac_q: Quantizer::new(dac_q_bit_count, dac_q_full_scale),
            ca_q: Quantizer::new(ca_q_bit_count, ca_q_full_scale),
            std_period_counter: 0,
            std_phase: 0.0,
            ref_period_counter: 0,
            ref_phase: 0.0,
            ref_jitter_phase: Vec::new(),
            dco_period_counter: 0,
            dco_phase: 0.0,
            dco_jitter_phase: Vec::new(),
            dco_jitter_record: Vec::new(),
            dco_time_record: Vec::new(),
            ref_jitter_record: Vec::new(),
            ref_time_record: Vec::new(),
            debug_record: Vec::new(),
            gain_record: Vec::new(),
            ref_jitter_acc: 0.0,
            dco_control_word: 0.0,
            dco_cw_raw: 0.0,
            ca_control_word: 0.0,
            ca_cw_raw: 0.0,
            sampler: 0.0,
            rawsampler: 0.0,
            last_sample: Vec::new(),
            delay_n,
        }
    }
    fn initial(&mut self) {
        self.dco_jitter_phase.truncate(0);
        self.ref_jitter_phase.truncate(0);
        self.dco_jitter_record.truncate(0);
        self.ref_jitter_record.truncate(0);
        self.dco_time_record.truncate(0);
        self.ref_time_record.truncate(0);
        self.debug_record.truncate(0);
        self.gain_record.truncate(0);
        self.ref_jitter_acc = self.rng_ref.sample(&mut thread_rng());
        self.dco_jitter_phase
            .push(self.rng_g.sample(&mut thread_rng()));
        self.ref_jitter_phase.push(self.ref_jitter_acc);
        self.dco_phase = 0.0;
        self.ref_phase = 0.0;
        self.std_phase = 0.0;
        self.std_period_counter = 0;
        self.dco_period_counter = 0;
        self.ref_period_counter = 0;
        self.lms_counter = 10.0;
        self.lms_counter_ca = 0.0;
        self.delta_sigma_counter = 0.0;
        self.alpha_counter = 0.0;
        (0..self.delay_n).for_each(|_| self.last_sample.push(0.0));
    }
    fn update_phase(&mut self) {
        let refdt = (1.0 - self.ref_phase) * (1.0 + self.ref_jitter_phase.last().unwrap());
        let stddt = 1.0 - self.std_phase;
        let dcodt = (1.0 - self.dco_phase) / (1.0 + self.dco_control_word * self.g_dac_vco)
            * (1.0 + self.dco_jitter_phase.last().unwrap());
        if (refdt < stddt) && (refdt < dcodt) {
            self.dco_phase += refdt * (1.0 + self.dco_control_word * self.g_dac_vco)
                / (1.0 + self.dco_jitter_phase.last().unwrap());
            self.ref_phase = 0.0;
            self.std_phase += refdt;
            self.ref_period_counter += 1;
            let t = self.rng_ref.sample(&mut thread_rng());
            self.ref_jitter_phase.push(t - self.ref_jitter_acc);
            self.ref_jitter_acc = t;
            self.ref_time_record
                .push(self.std_period_counter as f64 + self.std_phase);

            if self.ref_period_counter % ((self.n * 2) as u64) == 0 {
                self.gain_record.push(self.lms_counter_ca);
                let error = self.rawsampler
                    - self.lms_counter * self.delta_sigma_counter
                    - self.lms_counter_ca * self.delta_sigma_counter_ca;
                self.debug_record.push(error);
                self.lms_counter += self.lms_gain * self.delta_sigma_counter * error;
                self.lms_counter_ca += self.lms_gain_ca * self.delta_sigma_counter_ca * error;
                self.delta_sigma_counter +=
                    self.dco_cw_raw - self.delta_sigma_counter - self.dco_control_word;
                self.delta_sigma_counter_ca +=
                    self.ca_cw_raw - self.delta_sigma_counter_ca - self.ca_control_word;
                self.alpha_counter += self.sampler;
                self.rawsampler = if self.dco_period_counter % 2 == 1 {
                    1.0
                } else {
                    -1.0
                };
                if self.delay_n == 0 {
                    self.sampler = self.g * self.rawsampler * 2.2 / self.lms_counter;
                    // self.sampler = self.g * self.rawsampler;
                } else {
                    self.sampler = self.last_sample[0];
                    (0..self.delay_n - 1)
                        .for_each(|i| self.last_sample[i] = self.last_sample[i + 1]);
                    // self.last_sample[self.delay_n - 1] = self.g * 20.0 / self.lms_counter * self.rawsampler;
                    self.last_sample[self.delay_n - 1] = self.g * self.rawsampler
                }
                self.dco_cw_raw = self.alpha_counter + self.delta_sigma_counter;
                // self.ca_cw_raw = -self.g_charge_accelerating * self.rawsampler + self.delta_sigma_counter_ca;
                self.ca_cw_raw =
                    -self.g_charge_accelerating * self.rawsampler * 3.788 / self.lms_counter_ca + self.delta_sigma_counter_ca;
                self.dco_control_word = self.dac_q.quantify(self.dco_cw_raw);
                self.ca_control_word = self.ca_q.quantify(self.ca_cw_raw);
                self.debug_record.push(self.ca_control_word);
                self.dco_phase += self.ca_control_word;
            }
        } else if (dcodt < stddt) && (dcodt < refdt) {
            if dcodt >= 0.0 {
                self.dco_phase = 0.0;
                self.ref_phase += dcodt / (1.0 + self.ref_jitter_phase.last().unwrap());
                self.std_phase += dcodt;
                self.dco_period_counter += 1;
                self.dco_jitter_phase
                    .push(self.rng_g.sample(&mut thread_rng()));
                self.dco_time_record
                    .push(self.std_period_counter as f64 + self.std_phase);
            } else {
                self.dco_phase = self.dco_phase - 1.0;
                self.dco_period_counter += 1;
                self.dco_jitter_phase
                    .push(self.rng_g.sample(&mut thread_rng()));
                self.dco_time_record
                    .push(self.std_period_counter as f64 + self.std_phase);
            }
        } else {
            self.dco_phase += stddt * (1.0 + self.dco_control_word * self.g_dac_vco)
                / (1.0 + self.dco_jitter_phase.last().unwrap());
            self.ref_phase += stddt / (1.0 + self.ref_jitter_phase.last().unwrap());
            self.std_phase = 0.0;
            self.std_period_counter += 1;
            self.dco_jitter_record.push(if self.dco_phase > 0.5 {
                self.dco_phase - 1.0
            } else {
                self.dco_phase
            });
            self.ref_jitter_record.push(if self.ref_phase > 0.5 {
                self.ref_phase - 1.0
            } else {
                self.ref_phase
            });
        }
    }
    fn clear_record(&mut self) {
        self.dco_time_record.clear();
        self.dco_jitter_record.clear();
        self.ref_time_record.clear();
        self.ref_jitter_record.clear();
        self.debug_record.clear();
        self.gain_record.clear();
        if self.dco_jitter_phase.len() > 100 {
            (0..100).for_each(|index| {
                self.dco_jitter_phase[index] =
                    self.dco_jitter_phase[self.dco_jitter_phase.len() - 100 + index]
            });
            self.dco_jitter_phase.truncate(100);
        }
        if self.ref_jitter_phase.len() > 100 {
            (0..100).for_each(|index| {
                self.ref_jitter_phase[index] =
                    self.ref_jitter_phase[self.ref_jitter_phase.len() - 100 + index]
            });
            self.ref_jitter_phase.truncate(100);
        }
    }
}

fn save(path: String, data: &Vec<f64>) {
    let mut saver = File::create(path).unwrap();
    let mut data_string = String::new();
    data.iter()
        .for_each(|i| data_string += &(i.to_string() + "\n"));
    saver
        .write_all(data_string.as_bytes())
        .expect("write failed");
}

#[derive(Copy, Clone, Debug)]
struct GainMat {
    g_p: f64,
    g_a: f64,
}

fn main() {
    let start = SystemTime::now();
    let iterp = 1;
    let itera = 64;
    let prangemax = -0.0018;
    let prangemin = -0.0018;
    let arangemax = 0.03;
    let arangemin = 0.001;
    let rawiter = (0..(iterp + 1))
        .map(|t| {
            -t as f64 * prangemax / (iterp as f64) - (iterp - t) as f64 * prangemin / (iterp as f64)
        })
        .flat_map(|t| {
            (1..(itera + 1))
                .map(|t| {
                    -t as f64 * arangemax / (itera as f64)
                        - (itera - t) as f64 * arangemin / (itera as f64)
                })
                .map(move |t2| GainMat { g_p: t, g_a: t2 })
        })
        .collect::<Vec<GainMat>>();
    let data: Vec<Vec<f64>> = rawiter
        .par_iter()
        .map(|j| {
            let data = *j;
            let j = &data.g_p;
            let ja = &data.g_a;
            let begin: u64 = 1000000000;
            let end:   u64 = 10000000;

            let mut ans = vec![0., 0., 0., 0., 0., 0., 0.];
            let mut t = BBPLLState::new(
                0.003, 0.002, 5, *j, *ja, 5.0, 0.1, 0.015, 16, 0.2, 14, 1.0, 0,
            );
            t.initial();
            if begin > end {
                (0..((begin / end) + 1)).for_each(|index| {
                    (0..end).for_each(|_| {
                        if index <= 5 {
                            t.lms_counter = 3.0;
                            t.lms_counter_ca = 4.0;
                        }
                        t.update_phase()
                    });
                    t.clear_record();
                })
            } else {
                (0..begin).for_each(|_| t.update_phase());
                t.clear_record();
            }
            (0..end).for_each(|_| t.update_phase());
            let ref_mean =
                &t.ref_jitter_record.iter().sum::<f64>() / t.ref_jitter_record.len() as f64;
            let ref_std = (&t
                .ref_jitter_record
                .iter()
                .map(|x| (x - ref_mean) * (x - ref_mean))
                .sum::<f64>()
                / t.ref_jitter_record.len() as f64)
                .sqrt();
            let dco_mean =
                &t.dco_jitter_record.iter().sum::<f64>() / t.dco_jitter_record.len() as f64;
            let dco_std = (&t
                .dco_jitter_record
                .iter()
                .map(|x| (x - dco_mean) * (x - dco_mean))
                .sum::<f64>()
                / t.dco_jitter_record.len() as f64)
                .sqrt();
            let gain_mean = &t.gain_record.iter().sum::<f64>() / t.gain_record.len() as f64;
            // save(
            //     "/tmp/data/dcojitter".to_owned() + &(*j).to_string() + "_" + &(*ja).to_string() + ".dat",
            //     &t.dco_jitter_record,
            // );
            // save(
            //     "/tmp/data/dcotime".to_owned() + &(*j).to_string() + "_" + &(*ja).to_string() + ".dat",
            //     &t.dco_time_record,
            // );
            // save(
            //     "/tmp/data/refjitter".to_owned() + &(*j).to_string() + "_" + &(*ja).to_string() + ".dat",
            //     &t.ref_jitter_record,
            // );
            // save(
            //     "/tmp/data/reftime".to_owned() + &(*j).to_string() + "_" + &(*ja).to_string() + ".dat",
            //     &t.ref_time_record,
            // );
            save(
                "/tmp/data/debug".to_owned() + &(*j).to_string() + "_" + &(*ja).to_string() + ".dat",
                &t.debug_record,
            );
            save(
                "/tmp/data/gain".to_owned() + &(*j).to_string() + "_" + &(*ja).to_string() + ".dat",
                &t.gain_record,
            );
            ans[0] = *j as f64;
            ans[1] = *ja as f64;
            ans[2] = dco_mean;
            ans[3] = dco_std;
            ans[4] = ref_mean;
            ans[5] = ref_std;
            ans[6] = gain_mean;
            ans
        })
        .collect();
    println!(
        "Finish using: {:?}s",
        SystemTime::now().duration_since(start).unwrap().as_secs()
    );
    println!("gain dco, gain acc, jitter mean, dco jitter std, ref jitter mean, ref jitter std, gain mean");
    for i in data {
        println!(
            "{}, {}, {}, {}, {}, {}, {}",
            i[0], i[1], i[2], i[3], i[4], i[5], i[6]
        );
    }
}
